from config import conf
from os.path import join
import pandas as pd
import numpy as np
from data_utils.basic_data import load_basic_dataset
from functools import reduce

# ----------------------------------------------------

train_df = load_basic_dataset(split='train')
train_entities = set(reduce(lambda x, y: x + y, train_df['entity'].map(lambda x: str(x).split(';'))))
negative_entities = set(reduce(lambda x, y: x + y, train_df['key_entity'].map(lambda x: str(x).split(';'))))
nine = train_entities - negative_entities


# -----------------------------------------------------

def remove_short_entity(entity_str):
    """
    除去key_entity中同一实体的较短名称
    :param entity_str:
    :return:
    """
    if not isinstance(entity_str, str):
        return entity_str
    entities = entity_str.split(';')
    states = np.ones(len(entities))
    for i, e in enumerate(entities):
        for p in entities:
            if e in p and len(e) < len(p) and ('(' not in p) and ('（' not in p):
                states[i] = 0
                print('removed %s by %s' % (e, p))
    rs = []
    for i, e in enumerate(entities):
        if states[i] == 1:
            rs.append(e)
    return ';'.join(rs)


def remove_nine(key_entity):
    """
    除去不在训练集中negative entity的entity
    :param key_entity:
    :return:
    """
    if not isinstance(key_entity, str):
        return key_entity
    rs = []
    for e in key_entity.split(';'):
        if e in nine:
            print('removed %s by nine' % (e))
            continue
        else:
            rs.append(e)
    if len(rs) == 0:
        rs = np.nan
    else:
        rs = ';'.join(rs)
    return rs


if __name__ == '__main__':
    INFERENCE_DIR = conf.get('dir', 'inference_result_dir')
    RESULT_FILE = join(INFERENCE_DIR, 'evaluation', 'BertSentiEntityscore0.940010_epoch5.csv')
    df = pd.read_csv(RESULT_FILE)
    df['key_entity'] = df['key_entity'].map(remove_short_entity)
    df.to_csv(join(INFERENCE_DIR, 'evaluation', 'BertSentiEntityscore0.940010_epoch5_shortKeyRemoved_nineRemoved.csv'),
              index=False)
